package com.zhigh.project.security.test;

import com.google.common.collect.Streams;
import com.nimbusds.jose.JWSAlgorithm;
import com.zhigh.project.security.jwt.authentication.DefaultJWT;
import com.zhigh.project.security.jwt.authentication.DefaultJWTHeader;
import com.zhigh.project.security.jwt.authentication.DefaultJWTPayload;
import com.zhigh.project.security.jwt.authentication.JWT;
import com.zhigh.project.security.jwt.authentication.JWTAuthentication;
import com.zhigh.project.security.jwt.authentication.SecretHolder;
import com.zhigh.project.security.jwt.authentication.principal.DefaultJWTPrincipal;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.lang.NonNull;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.test.context.support.WithSecurityContextFactory;

import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

public class WithMockJWTSecurityContextFactory implements WithSecurityContextFactory<WithMockJWT>, ApplicationContextAware {

    private ApplicationContext context;

    private final SecretHolder<JWT> secretHolder;

    public WithMockJWTSecurityContextFactory(SecretHolder<JWT> secretHolder) {
        this.secretHolder = secretHolder;
    }

    @Override
    public SecurityContext createSecurityContext(WithMockJWT annotation) {
        String jti = UUID.randomUUID().toString();

        List<String> authorities = Streams.concat(Arrays.stream(annotation.authorities()), Arrays.stream(annotation.roles()).map(role -> "ROLE_" + role)).collect(Collectors.toList());
        DefaultJWT jwt = new DefaultJWT(
                new DefaultJWTHeader(JWSAlgorithm.HS256.getName()),
                DefaultJWTPayload.builder().username(annotation.username()).authorities(authorities).jti(jti).build(),
                jti
        );

        JWTAuthentication authentication = new JWTAuthentication(new DefaultJWTPrincipal(jwt));
        secretHolder.saveSecret(jwt, jti);
        authentication.setDetails(jwt);

        SecurityContext context = SecurityContextHolder.createEmptyContext();
        context.setAuthentication(authentication);
        return context;
    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext context) throws BeansException {
        this.context = context;
    }
}
